import asyncio
import os

from agents import (
    Agent,
    Model,
    ModelProvider,
    OpenAIChatCompletionsModel,
    RunConfig,
    Runner,
    function_tool,
    set_tracing_disabled,
)
from dotenv import load_dotenv
from openai import AsyncOpenAI

# 环境变量加载相关

# 加载.env文件
load_dotenv()
# 从.env文件中获取BASE_URL、API_KEY、MODEL_NAME信息
BASE_URL = os.getenv("BASE_URL") or ""
API_KEY = os.getenv("API_KEY") or ""
MODEL_NAME = os.getenv("MODEL_NAME") or ""

if not BASE_URL or not API_KEY or not MODEL_NAME:
    raise ValueError(
        "请通过环境变量或代码设置EXAMPLE_BASE_URL、EXAMPLE_API_KEY、EXAMPLE_MODEL_NAME。"
    )

"""
本例使用自定义提供程序调用Runner.run()的部分，并直接调用OpenAI进行其他操作。
步骤：1. 创建一个自定义的OpenAI客户端。
2. 创建一个使用自定义客户端的ModelProvider。
3. 在调用Runner.run()时使用ModelProvider，仅在我们想使用自定义LLM提供程序时。
注意，在本例中，我们假设您没有从platform.openai.com获取API密钥，因此禁用了跟踪。
如果您有API密钥，您可以选择设置`OPENAI_API_KEY`环境变量或调用set_tracing_export_api_key()来设置跟踪特定的密钥。
"""
# ①创建一个自定义的OpenAI客户端-Client
client = AsyncOpenAI(base_url=BASE_URL, api_key=API_KEY)
set_tracing_disabled(disabled=True)


# ②创建一个使用自定义客户端的-ModelProvider
class DeepSeekModelProvider(ModelProvider):
    def get_model(self, model_name: str | None) -> Model:
        return OpenAIChatCompletionsModel(model=model_name or MODEL_NAME, openai_client=client)


# ③实例化ModelProvider对象
CUSTOM_MODEL_PROVIDER = DeepSeekModelProvider()


# ④定义一个工具函数
@function_tool
def get_weather(city: str):
    print(f"[debug] getting weather for {city}")
    return f"The weather in {city} is sunny."


async def main():
    # ⑤创建一个Agent
    agent = Agent(name="Assistant", instructions="你只能用中文进行回复", tools=[get_weather])

    # ⑥获取Agent对话的结果
    result = await Runner.run(
        agent,
        "给我讲一个笑话吧！",
        run_config=RunConfig(model_provider=CUSTOM_MODEL_PROVIDER),
    )
    print(result.final_output)

    # ⑦获取Agent对话的结果-没有配置RunConfig使用的是OpenAI的默认模型
    # result = await Runner.run(
    #     agent,
    #     "给我讲一个笑话吧!",
    # )
    # print(result.final_output)


if __name__ == "__main__":
    asyncio.run(main())
